{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| default_exp models.deepar"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DeepAR"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The DeepAR model produces probabilistic forecasts based on an autoregressive recurrent neural network optimized on panel data using cross-learning. DeepAR obtains its forecast distribution uses a Markov Chain Monte Carlo sampler with the following conditional probability:\n",
    "$$\\mathbb{P}(\\mathbf{y}_{[t+1:t+H]}|\\;\\mathbf{y}_{[:t]},\\; \\mathbf{x}^{(f)}_{[:t+H]},\\; \\mathbf{x}^{(s)})$$\n",
    "\n",
    "where $\\mathbf{x}^{(s)}$ are static exogenous inputs, $\\mathbf{x}^{(f)}_{[:t+H]}$ are future exogenous available at the time of the prediction.\n",
    "The predictions are obtained by transforming the hidden states $\\mathbf{h}_{t}$ into predictive distribution parameters $\\theta_{t}$, and then generating samples $\\mathbf{\\hat{y}}_{[t+1:t+H]}$ through Monte Carlo sampling trajectories.\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "\\mathbf{h}_{t} &= \\textrm{RNN}([\\mathbf{y}_{t},\\mathbf{x}^{(f)}_{t+1},\\mathbf{x}^{(s)}], \\mathbf{h}_{t-1})\\\\\n",
    "\\mathbf{\\theta}_{t}&=\\textrm{Linear}(\\mathbf{h}_{t}) \\\\\n",
    "\\hat{y}_{t+1}&=\\textrm{sample}(\\;\\mathrm{P}(y_{t+1}\\;|\\;\\mathbf{\\theta}_{t})\\;)\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "**References**<br>\n",
    "- [David Salinas, Valentin Flunkert, Jan Gasthaus, Tim Januschowski (2020). \"DeepAR: Probabilistic forecasting with autoregressive recurrent networks\". International Journal of Forecasting.](https://www.sciencedirect.com/science/article/pii/S0169207019301888)<br>\n",
    "- [Alexander Alexandrov et. al (2020). \"GluonTS: Probabilistic and Neural Time Series Modeling in Python\". Journal of Machine Learning Research.](https://www.jmlr.org/papers/v21/19-820.html)<br>\n",
    "\n",
    "\n",
    ":::{.callout-warning collapse=\"false\"}\n",
    "#### Exogenous Variables, Losses, and Parameters Availability\n",
    "\n",
    "Given the sampling procedure during inference, DeepAR only supports `DistributionLoss` as training loss.\n",
    "\n",
    "Note that DeepAR generates a non-parametric forecast distribution using Monte Carlo. We use this sampling procedure also during validation to make it closer to the inference procedure. Therefore, only the `MQLoss` is available for validation.\n",
    "\n",
    "Aditionally, Monte Carlo implies that historic exogenous variables are not available for the model.\n",
    ":::"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Figure 1. DeepAR model, during training the optimization signal comes from likelihood of observations, during inference a recurrent multi-step strategy is used to generate predictive distributions.](imgs_models/deepar.jpeg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "from typing import Optional\n",
    "\n",
    "from neuralforecast.common._base_model import BaseModel\n",
    "from neuralforecast.losses.pytorch import DistributionLoss, MAE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "import logging\n",
    "import warnings\n",
    "\n",
    "from fastcore.test import test_eq\n",
    "from nbdev.showdoc import show_doc\n",
    "from neuralforecast.common._model_checks import check_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "logging.getLogger(\"pytorch_lightning\").setLevel(logging.ERROR)\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class Decoder(nn.Module):\n",
    "    \"\"\"Multi-Layer Perceptron Decoder\n",
    "\n",
    "    **Parameters:**<br>\n",
    "    `in_features`: int, dimension of input.<br>\n",
    "    `out_features`: int, dimension of output.<br>\n",
    "    `hidden_size`: int, dimension of hidden layers.<br>\n",
    "    `num_layers`: int, number of hidden layers.<br>\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, in_features, out_features, hidden_size, hidden_layers):\n",
    "        super().__init__()\n",
    "\n",
    "        if hidden_layers == 0:\n",
    "            # Input layer\n",
    "            layers = [nn.Linear(in_features=in_features, out_features=out_features)]\n",
    "        else:\n",
    "            # Input layer\n",
    "            layers = [nn.Linear(in_features=in_features, out_features=hidden_size), nn.ReLU()]\n",
    "            # Hidden layers\n",
    "            for i in range(hidden_layers - 2):\n",
    "                layers += [nn.Linear(in_features=hidden_size, out_features=hidden_size), nn.ReLU()]\n",
    "            # Output layer\n",
    "            layers += [nn.Linear(in_features=hidden_size, out_features=out_features)]\n",
    "\n",
    "        # Store in layers as ModuleList\n",
    "        self.layers = nn.Sequential(*layers)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.layers(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class DeepAR(BaseModel):\n",
    "    \"\"\" DeepAR\n",
    "\n",
    "    **Parameters:**<br>\n",
    "    `h`: int, Forecast horizon. <br>\n",
    "    `input_size`: int, maximum sequence length for truncated train backpropagation. Default -1 uses 3 * horizon <br>\n",
    "    `h_train`: int, maximum sequence length for truncated train backpropagation. Default 1.<br>\n",
    "    `lstm_n_layers`: int=2, number of LSTM layers.<br>\n",
    "    `lstm_hidden_size`: int=128, LSTM hidden size.<br>\n",
    "    `lstm_dropout`: float=0.1, LSTM dropout.<br>\n",
    "    `decoder_hidden_layers`: int=0, number of decoder MLP hidden layers. Default: 0 for linear layer. <br>\n",
    "    `decoder_hidden_size`: int=0, decoder MLP hidden size. Default: 0 for linear layer.<br>\n",
    "    `trajectory_samples`: int=100, number of Monte Carlo trajectories during inference.<br>\n",
    "    `stat_exog_list`: str list, static exogenous columns.<br>\n",
    "    `hist_exog_list`: str list, historic exogenous columns.<br>\n",
    "    `futr_exog_list`: str list, future exogenous columns.<br>\n",
    "    `exclude_insample_y`: bool=False, the model skips the autoregressive features y[t-input_size:t] if True.<br>\n",
    "    `loss`: PyTorch module, instantiated train loss class from [losses collection](https://nixtla.github.io/neuralforecast/losses.pytorch.html).<br>\n",
    "    `valid_loss`: PyTorch module=`loss`, instantiated valid loss class from [losses collection](https://nixtla.github.io/neuralforecast/losses.pytorch.html).<br>\n",
    "    `max_steps`: int=1000, maximum number of training steps.<br>\n",
    "    `learning_rate`: float=1e-3, Learning rate between (0, 1).<br>\n",
    "    `num_lr_decays`: int=-1, Number of learning rate decays, evenly distributed across max_steps.<br>\n",
    "    `early_stop_patience_steps`: int=-1, Number of validation iterations before early stopping.<br>\n",
    "    `val_check_steps`: int=100, Number of training steps between every validation loss check.<br>\n",
    "    `batch_size`: int=32, number of different series in each batch.<br>\n",
    "    `valid_batch_size`: int=None, number of different series in each validation and test batch, if None uses batch_size.<br>\n",
    "    `windows_batch_size`: int=1024, number of windows to sample in each training batch, default uses all.<br>\n",
    "    `inference_windows_batch_size`: int=-1, number of windows to sample in each inference batch, -1 uses all.<br>\n",
    "    `start_padding_enabled`: bool=False, if True, the model will pad the time series with zeros at the beginning, by input size.<br>\n",
    "    `step_size`: int=1, step size between each window of temporal data.<br>\n",
    "    `scaler_type`: str='identity', type of scaler for temporal inputs normalization see [temporal scalers](https://nixtla.github.io/neuralforecast/common.scalers.html).<br>\n",
    "    `random_seed`: int, random_seed for pytorch initializer and numpy generators.<br>\n",
    "    `drop_last_loader`: bool=False, if True `TimeSeriesDataLoader` drops last non-full batch.<br>\n",
    "    `alias`: str, optional,  Custom name of the model.<br>\n",
    "    `optimizer`: Subclass of 'torch.optim.Optimizer', optional, user specified optimizer instead of the default choice (Adam).<br>\n",
    "    `optimizer_kwargs`: dict, optional, list of parameters used by the user specified `optimizer`.<br>\n",
    "    `lr_scheduler`: Subclass of 'torch.optim.lr_scheduler.LRScheduler', optional, user specified lr_scheduler instead of the default choice (StepLR).<br>\n",
    "    `lr_scheduler_kwargs`: dict, optional, list of parameters used by the user specified `lr_scheduler`.<br>\n",
    "    `dataloader_kwargs`: dict, optional, list of parameters passed into the PyTorch Lightning dataloader by the `TimeSeriesDataLoader`. <br>\n",
    "    `**trainer_kwargs`: int,  keyword trainer arguments inherited from [PyTorch Lighning's trainer](https://pytorch-lightning.readthedocs.io/en/stable/api/pytorch_lightning.trainer.trainer.Trainer.html?highlight=trainer).<br>    \n",
    "\n",
    "    **References**<br>\n",
    "    - [David Salinas, Valentin Flunkert, Jan Gasthaus, Tim Januschowski (2020). \"DeepAR: Probabilistic forecasting with autoregressive recurrent networks\". International Journal of Forecasting.](https://www.sciencedirect.com/science/article/pii/S0169207019301888)<br>\n",
    "    - [Alexander Alexandrov et. al (2020). \"GluonTS: Probabilistic and Neural Time Series Modeling in Python\". Journal of Machine Learning Research.](https://www.jmlr.org/papers/v21/19-820.html)<br>\n",
    "\n",
    "    \"\"\"\n",
    "    # Class attributes\n",
    "    EXOGENOUS_FUTR = True\n",
    "    EXOGENOUS_HIST = False\n",
    "    EXOGENOUS_STAT = True\n",
    "    MULTIVARIATE = False\n",
    "    RECURRENT = True\n",
    "\n",
    "    def __init__(self,\n",
    "                 h,\n",
    "                 input_size: int = -1,\n",
    "                 h_train: int = 1,\n",
    "                 lstm_n_layers: int = 2,\n",
    "                 lstm_hidden_size: int = 128,\n",
    "                 lstm_dropout: float = 0.1,\n",
    "                 decoder_hidden_layers: int = 0,\n",
    "                 decoder_hidden_size: int = 0,\n",
    "                 trajectory_samples: int = 100,\n",
    "                 stat_exog_list = None,\n",
    "                 hist_exog_list = None,\n",
    "                 futr_exog_list = None,\n",
    "                 exclude_insample_y = False,\n",
    "                 loss = DistributionLoss(distribution='StudentT', level=[80, 90], return_params=False),\n",
    "                 valid_loss = MAE(),\n",
    "                 max_steps: int = 1000,\n",
    "                 learning_rate: float = 1e-3,\n",
    "                 num_lr_decays: int = 3,\n",
    "                 early_stop_patience_steps: int =-1,\n",
    "                 val_check_steps: int = 100,\n",
    "                 batch_size: int = 32,\n",
    "                 valid_batch_size: Optional[int] = None,\n",
    "                 windows_batch_size: int = 1024,\n",
    "                 inference_windows_batch_size: int = -1,\n",
    "                 start_padding_enabled = False,\n",
    "                 step_size: int = 1,\n",
    "                 scaler_type: str = 'identity',\n",
    "                 random_seed: int = 1,\n",
    "                 drop_last_loader = False,\n",
    "                 alias: Optional[str] = None,\n",
    "                 optimizer = None,\n",
    "                 optimizer_kwargs = None,\n",
    "                 lr_scheduler = None,\n",
    "                 lr_scheduler_kwargs = None,\n",
    "                 dataloader_kwargs = None,\n",
    "                 **trainer_kwargs):\n",
    "\n",
    "        if exclude_insample_y:\n",
    "            raise Exception('DeepAR has no possibility for excluding y.')\n",
    "        \n",
    "        # Inherit BaseWindows class\n",
    "        super(DeepAR, self).__init__(h=h,\n",
    "                                    input_size=input_size,\n",
    "                                    h_train=h_train,\n",
    "                                    stat_exog_list=stat_exog_list,\n",
    "                                    hist_exog_list=hist_exog_list,\n",
    "                                    futr_exog_list=futr_exog_list,\n",
    "                                    exclude_insample_y = exclude_insample_y,\n",
    "                                    loss=loss,\n",
    "                                    valid_loss=valid_loss,\n",
    "                                    max_steps=max_steps,\n",
    "                                    learning_rate=learning_rate,\n",
    "                                    num_lr_decays=num_lr_decays,\n",
    "                                    early_stop_patience_steps=early_stop_patience_steps,\n",
    "                                    val_check_steps=val_check_steps,\n",
    "                                    batch_size=batch_size,\n",
    "                                    valid_batch_size=valid_batch_size,\n",
    "                                    windows_batch_size=windows_batch_size,\n",
    "                                    inference_windows_batch_size=inference_windows_batch_size,\n",
    "                                    start_padding_enabled=start_padding_enabled,\n",
    "                                    step_size=step_size,\n",
    "                                    scaler_type=scaler_type,\n",
    "                                    random_seed=random_seed,\n",
    "                                    drop_last_loader=drop_last_loader,\n",
    "                                    alias=alias,\n",
    "                                    optimizer=optimizer,\n",
    "                                    optimizer_kwargs=optimizer_kwargs,\n",
    "                                    lr_scheduler=lr_scheduler,\n",
    "                                    lr_scheduler_kwargs=lr_scheduler_kwargs,\n",
    "                                    dataloader_kwargs=dataloader_kwargs,\n",
    "                                    **trainer_kwargs)\n",
    "\n",
    "        self.n_samples = trajectory_samples\n",
    "\n",
    "        # LSTM\n",
    "        self.encoder_n_layers = lstm_n_layers\n",
    "        self.encoder_hidden_size = lstm_hidden_size\n",
    "        self.encoder_dropout = lstm_dropout\n",
    "       \n",
    "        # LSTM input size (1 for target variable y)\n",
    "        input_encoder = 1 + self.futr_exog_size + self.stat_exog_size\n",
    "\n",
    "        # Instantiate model\n",
    "        self.rnn_state = None\n",
    "        self.maintain_state = False\n",
    "        self.hist_encoder = nn.LSTM(input_size=input_encoder,\n",
    "                                    hidden_size=self.encoder_hidden_size,\n",
    "                                    num_layers=self.encoder_n_layers,\n",
    "                                    dropout=self.encoder_dropout,\n",
    "                                    batch_first=True)\n",
    "\n",
    "        # Decoder MLP\n",
    "        self.decoder = Decoder(in_features=lstm_hidden_size,\n",
    "                               out_features=self.loss.outputsize_multiplier,\n",
    "                               hidden_size=decoder_hidden_size,\n",
    "                               hidden_layers=decoder_hidden_layers)\n",
    "\n",
    "    def forward(self, windows_batch):\n",
    "\n",
    "        # Parse windows_batch\n",
    "        encoder_input = windows_batch['insample_y'] # <- [B, T, 1]\n",
    "        futr_exog  = windows_batch['futr_exog']\n",
    "        stat_exog  = windows_batch['stat_exog']\n",
    "\n",
    "        _, input_size = encoder_input.shape[:2]\n",
    "        if self.futr_exog_size > 0:\n",
    "            encoder_input = torch.cat((encoder_input, futr_exog), dim=2)\n",
    "\n",
    "        if self.stat_exog_size > 0:\n",
    "            stat_exog = stat_exog.unsqueeze(1).repeat(1, input_size, 1)     # [B, S] -> [B, input_size-1, S]\n",
    "            encoder_input = torch.cat((encoder_input, stat_exog), dim=2)\n",
    "\n",
    "        # RNN forward\n",
    "        if self.maintain_state:\n",
    "            rnn_state = self.rnn_state\n",
    "        else:\n",
    "            rnn_state = None\n",
    "\n",
    "        hidden_state, rnn_state = self.hist_encoder(encoder_input, \n",
    "                                                    rnn_state)              # [B, input_size-1, rnn_hidden_state]\n",
    "\n",
    "        if self.maintain_state:\n",
    "            self.rnn_state = rnn_state\n",
    "\n",
    "        # Decoder forward\n",
    "        output = self.decoder(hidden_state)                                 # [B, input_size-1, output_size]\n",
    "\n",
    "        # Return only horizon part\n",
    "        return output[:, -self.h:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(DeepAR, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(DeepAR.fit, name='DeepAR.fit', title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(DeepAR.predict, name='DeepAR.predict', title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Unit tests for models\n",
    "logging.getLogger(\"pytorch_lightning\").setLevel(logging.ERROR)\n",
    "logging.getLogger(\"lightning_fabric\").setLevel(logging.ERROR)\n",
    "with warnings.catch_warnings():\n",
    "    warnings.simplefilter(\"ignore\")\n",
    "    check_model(DeepAR, [\"airpassengers\"])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Usage Example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| eval: false\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from neuralforecast import NeuralForecast\n",
    "from neuralforecast.models import DeepAR\n",
    "from neuralforecast.losses.pytorch import DistributionLoss, MQLoss\n",
    "from neuralforecast.utils import AirPassengersPanel, AirPassengersStatic\n",
    "Y_train_df = AirPassengersPanel[AirPassengersPanel.ds<AirPassengersPanel['ds'].values[-12]] # 132 train\n",
    "Y_test_df = AirPassengersPanel[AirPassengersPanel.ds>=AirPassengersPanel['ds'].values[-12]].reset_index(drop=True) # 12 test\n",
    "\n",
    "nf = NeuralForecast(\n",
    "    models=[DeepAR(h=12,\n",
    "                   input_size=24,\n",
    "                   lstm_n_layers=1,\n",
    "                   trajectory_samples=100,\n",
    "                   loss=DistributionLoss(distribution='StudentT', level=[80, 90], return_params=True),\n",
    "                   valid_loss=MQLoss(level=[80, 90]),\n",
    "                   learning_rate=0.005,\n",
    "                   stat_exog_list=['airline1'],\n",
    "                   futr_exog_list=['trend'],\n",
    "                   max_steps=100,\n",
    "                   val_check_steps=10,\n",
    "                   early_stop_patience_steps=-1,\n",
    "                   scaler_type='standard',\n",
    "                   enable_progress_bar=True,\n",
    "                   ),\n",
    "    ],\n",
    "    freq='ME'\n",
    ")\n",
    "nf.fit(df=Y_train_df, static_df=AirPassengersStatic, val_size=12)\n",
    "Y_hat_df = nf.predict(futr_df=Y_test_df)\n",
    "\n",
    "# Plot quantile predictions\n",
    "Y_hat_df = Y_hat_df.reset_index(drop=False).drop(columns=['unique_id','ds'])\n",
    "plot_df = pd.concat([Y_test_df, Y_hat_df], axis=1)\n",
    "plot_df = pd.concat([Y_train_df, plot_df])\n",
    "\n",
    "plot_df = plot_df[plot_df.unique_id=='Airline1'].drop('unique_id', axis=1)\n",
    "plt.plot(plot_df['ds'], plot_df['y'], c='black', label='True')\n",
    "plt.plot(plot_df['ds'], plot_df['DeepAR-median'], c='blue', label='median')\n",
    "plt.fill_between(x=plot_df['ds'][-12:], \n",
    "                 y1=plot_df['DeepAR-lo-90'][-12:].values, \n",
    "                 y2=plot_df['DeepAR-hi-90'][-12:].values,\n",
    "                 alpha=0.4, label='level 90')\n",
    "plt.legend()\n",
    "plt.grid()\n",
    "plt.plot()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
